module butcher_tableau_mod
  
  use constants_mod
  use namelist_mod

  implicit none

  private
  public :: butcher_tableau_init
  public :: butcher_tableau_final
  public :: nstage
  public :: IRK_a, IRK_b
  public :: ERK_a, ERK_b

  real(r8), allocatable :: ERK_a(:,:), ERK_b(:)
  real(r8), allocatable :: IRK_a(:,:), IRK_b(:)
  integer :: nstage

contains

  subroutine butcher_tableau_init()
    
    real(r8), parameter :: gamma = 1.0_r8 - 1.0_r8 / sqrt(2.0_r8)
    real(r8), parameter :: alpha = 0.2416942608_r8, &
                            beta = 0.0604235652_r8, &
                             eta = 0.1291528696_r8, &
                           delta = 0.5_r8 - beta - eta - alpha
    
    call butcher_tableau_final()

    select case (trim(time_scheme))
    case ('ERK_FE') ! Forward Euler
      nstage = 1
      allocate(ERK_a(nstage,nstage))
      allocate(ERK_b(nstage))
      ERK_a(1,1) = 0.0_r8
      ERK_b(1) = 1.0_r8
    case ('ERK3_WRF')
      nstage = 3
      allocate(ERK_a(nstage,nstage))
      allocate(ERK_b(nstage))
      ERK_a(1, 1) = 0.0_r8
      ERK_a(2, 1) = 1.0_r8 / 3.0_r8
      ERK_a(3, 1:2) = [0.0_r8, 0.5_r8]
      ERK_b(1:3) = [0.0_r8, 0.0_r8, 1.0_r8]
    case ('ERK4')
      nstage = 4
      !  0  | 0
      ! 1/2 | 1/2 
      ! 1/2 | 0   1/2  
      !  1  | 0   0    1   
      !  ---------------------
      !     | 1/6 1/3  1/3  1/6
      allocate(ERK_a(nstage,nstage))
      allocate(ERK_b(nstage))
      ! ERK_c(1:4  ) = [0.0_r8, 0.5_r8, 0.5_r8, 1.0_r8]
      ERK_a(1,1  ) = 0.0_r8
      ERK_a(2,1  ) = 0.5_r8
      ERK_a(3,1:2) = [0.0_r8, 0.5_r8]
      ERK_a(4,1:3) = [0.0_r8, 0.0_r8, 1.0_r8]
      ERK_b(1:4)   = [1.0_r8/6.0_r8, 1.0_r8/3.0_r8, 1.0_r8/3.0_r8, 1.0_r8/6.0_r8]
    case ('Backward_Euler')
      !  1 | 1 
      !  ------
      !    | 1 
      nstage = 1
      allocate(IRK_a(nstage,nstage))
      allocate(IRK_b(nstage))
      ! if (.not. allocated(RK_c)) allocate(RK_c(nstage))
      ! RK_c(1)   = 1.0_r8
      IRK_a(1,1) = 1.0_r8
      IRK_b(1)   = 1.0_r8
    case ('implicit_midpoint')
      !  1/2 | 1/2
      !  -------
      !    | 1
      nstage = 1
      allocate(IRK_a(nstage,nstage))
      allocate(IRK_b(nstage))
      ! if (.not. allocated(RK_c)) allocate(RK_c(nstage))
      ! RK_c(1)   = 0.5_r8
      IRK_a(1,1) = 0.5_r8
      IRK_b(1)   = 1.0_r8
    case ('Crank_Nicolson')
      !  Crank-Nicolson (implicit trapezoidal rule)
      ! 0 | 0   0
      ! 1 | 1/2 1/2
      !--------------
      !   | 1/2 1/2  
      nstage = 2
      allocate(IRK_a(nstage,nstage))
      allocate(IRK_b(nstage))
      ! if (.not. allocated(RK_c)) allocate(RK_c(nstage))
      ! RK_c(1:2)   = [0.0_r8, 1.0_r8]
      IRK_a(1,1:2) = [0.0_r8, 0.0_r8]
      IRK_a(2,1:2) = [0.5_r8, 0.5_r8]
      IRK_b = [0.5_r8, 0.5_r8]
    case ('ssp2') ! Weller 2013
      nstage = 2
      allocate(IRK_a(nstage,nstage))
      allocate(IRK_b(nstage))
      ! if (.not. allocated(RK_c)) allocate(RK_c(nstage))
      ! RK_c(1:2)   = [gamma, 1.0_r8 - gamma]
      IRK_a(1,1:2) = [gamma, 0.0_r8]
      IRK_a(2,1:2) = [1.0_r8 - 2*gamma, gamma] 
      IRK_b = [0.5_r8, 0.5_r8]
    case ('ssp3') ! weller 2013 ssp3(4,3,3), PR05
      nstage = 4
      allocate(IRK_a(nstage,nstage))
      allocate(IRK_b(nstage))
      IRK_a(1,1  ) = alpha
      IRK_a(2,1:2) = [-alpha, alpha]
      IRK_a(3,1:3) = [0.0_r8, 1.0_r8-alpha, alpha]
      IRK_a(4,1:4) = [beta, eta, delta, alpha]
      IRK_b(1:4) = [0.0_r8, 1.0_r8/6.0_r8, 1.0_r8/6.0_r8, 2.0_r8/3.0_r8]
    case default
      print*, "Unknown time_scheme in 'butcher_tableau_init'", trim(time_scheme)
      stop
    end select

  end subroutine butcher_tableau_init

  subroutine butcher_tableau_final

    if (allocated(IRK_a)) deallocate(IRK_a)
    if (allocated(IRK_b)) deallocate(IRK_b)
    if (allocated(ERK_a)) deallocate(ERK_a)
    if (allocated(ERK_b)) deallocate(ERK_b)

  end subroutine butcher_tableau_final

end module butcher_tableau_mod